from lib import *

def CHW2MK(ifm, layer):
    # 输出特征图
    out_height = layer['out_height']
    out_width = layer['out_width']
    # 输入特征图
    in_height = layer['in_height']
    in_width = layer['in_width']
    in_channel = layer['in_channel']
    # 卷积核大小
    kernel_h = layer['kernel_h']
    kernel_w = layer['kernel_w']
    # 控制卷积操作中的步长
    stride_h = layer['stride_h']
    stride_w = layer['stride_w']
    # 控制对输入特征图边缘填充像素的个数
    pad_h = layer['pad_h']
    pad_w = layer['pad_w']
    # 卷积核内部像素之间的步长
    dilation_h = layer['dilation_h']
    dilation_w = layer['dilation_w']
    # 输出特征图由输入特征图拉长而来，依照这样来看就是需要权重进行滑动
    # 这一行的数就是由卷积核滑动 Ci Hi Wi 得来的？（但貌似处理器比较少，可能得分两次滑才能得到最终结果）？
    M = out_height * out_width
    K = in_channel * kernel_h * kernel_w

    ofm = np.zeros((M, K))
    for oh in range(out_height):
        for ow in range(out_width):
            for kh in range(kernel_h):
                for kw in range(kernel_w):
                    ih = oh * stride_h - pad_h + kh * dilation_h
                    iw = ow * stride_w - pad_w + kw *dilation_w
                    if (0 <= ih < in_height) and (0 <= iw < in_width):
                        for ic in range(in_channel):
                            ofm[oh*out_width + ow][ic*kernel_h*kernel_w+kh*kernel_w+kw] = ifm[ic][ih][iw]
    return ofm

def test_im2col():
    in_height = np.random.randint(3, 100)
    in_width = np.random.randint(3, 100)
    in_channel = np.random.randint(3, 100)
    out_channel = np.random.randint(3, 100)
    kernel_h = np.random.randint(1,7)
    kernel_w = np.random.randint(1,7)
    stride_h = np.random.randint(1,7)
    stride_w = np.random.randint(1,7)
    pad_h = np.random.randint(0, kernel_h)
    pad_w = np.random.randint(0, kernel_w)
    dilation_h = np.random.randint(1,7)
    dilation_w = np.random.randint(1,7)
    # 计算公式不必说
    out_height = (in_height + 2 * pad_h - dilation_h * (kernel_h - 1) - 1) // stride_h + 1
    out_width = (in_width + 2 * pad_w - dilation_w * (kernel_w - 1) - 1) // stride_w + 1
    if(out_height <= 0 or out_width <= 0):
        return False
    # 定义卷积层参数
    layer = {
    'in_height': in_height,
    'in_width': in_width,
    'out_height': out_height,
    'out_width': out_width,
    'in_channel': in_channel,
    'out_channel': out_channel,
    'kernel_h': kernel_h,
    'kernel_w': kernel_w,
    'stride_h': stride_h,
    'stride_w': stride_w,
    'pad_h': pad_h,
    'pad_w': pad_w,
    'dilation_h': dilation_h,
    'dilation_w': dilation_w
    }
    print(layer)

    # 创建输入数据
    ifm = np.arange(in_channel * in_height * in_width).reshape(in_channel, in_height, in_width) + 1
    weight = np.arange(out_channel * in_channel * kernel_h * kernel_w).reshape(out_channel, in_channel, kernel_h, kernel_w) + 1
    bias = np.arange(out_channel) + 1
    # 将输入特征图转成矩阵 理论上的话 MK里同一个k共用一个权重，因为是要滑动的，所以在同一行？
    # 依据其 K = in_channel*kernel_h*kernel 相当于是把卷积核权重给拉直变成一行M数据了
    # MK中一行就是权重卷积核在输入特征图上滑动一次所要做计算的值
    # ifm_mk相当于将输入特征图数据拉成MK形式，
    ifm_mk = CHW2MK(ifm, layer)
    # 也就是说着所谓的C0通道是指不同的权重，就是out_channel 每个都对应不同的权重
    weight_nk = weight.reshape((out_channel, in_channel*kernel_h*kernel_w))
    ofm_mn = np.matmul(ifm_mk, weight_nk.transpose())
    ofm_nm = ofm_mn.transpose()
    for i in range(out_channel):
        ofm_nm[i] += bias[i]
    ofm_golden = golden_convolution(ifm, weight, bias, layer)
    compare(ofm_nm.flatten(), ofm_golden.flatten())
    return ofm_nm

for i in range(10):
    test_im2col()
